"""
Модуль для визуализации объектов и кластеров в 3D-графиках.

Содержит функции для создания 3D-графиков, включая объекты, кластеры и связи между ними.

Функции:
- `print_objects_and_clusters`: Визуализирует объекты и кластеры.
- `print_objects`: Визуализирует объекты в 3D-графике.
- `print_clusters`: Визуализирует кластеры объектов на 3D-графике.
- `print_leaders`: Визуализирует главные объекты и маршруты связи на 3D-графике.
- `print_short_paths`: Визуализирует кластеры объектов, маршруты связи и базовую станцию на 3D-графике.
- `print_graph_connections`: Визуализирует связи графа, включая базовую станцию и объекты.
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import networkx as nx
from swarm_network import data

# Чтение конфигурации и установка параметров генерации
config = data.read_config('config.yaml')

# Переводим размер в дюймы (1 дюйм = 80 пикселей при разрешении 80 dpi)
width_inches = config['fig_size_width'] / config['display_resolution']
height_inches = config['fig_size_height'] / config['display_resolution']

FIG_SIZE = (width_inches,height_inches)
FONT_SIZE = config['font_size']
FONT_SIZE_MEDIUM = config['font_size_medium']
STEP_X = config['text_step_x']
STEP_Y = config['text_step_y']
STEP_Z = config['text_step_z']

def print_objects_and_clusters(df_objects, df_clusters, base_point = None, print_id = True):
    """
    Визуализирует объекты и кластеры в 3D-графиках.

    Args:
        df_objects (pd.DataFrame): Данные об объектах.
        df_clusters (pd.DataFrame): Данные о кластерах.
        base_point (tuple, optional): Координаты базовой точки. По умолчанию None.
        print_id (bool, optional): Выводить идентификаторы объектов. По умолчанию True.
    """
    left, right = _prepare_axs()

    left = _prepare_objects(left, df_objects, print_id)

    if base_point is not None:
        left = _prepare_base(left, base_point)

    max_z = int(np.ceil(df_objects['z'].max() * 1.1 / 100) * 100)
    left = _ax_end(left, max_z, "Визуализация объектов", 1, 8, 'lower center')

    right = _prepare_clusters(right, df_objects, df_clusters)

    if base_point is not None:
        right = _prepare_base(right, base_point)

    right = _add_objects_ids(right, df_objects, 'black')

    max_z = int(np.ceil(df_objects['z'].max() * 1.1 / 100) * 100)
    right = _ax_end(right, max_z, "Визуализация кластеров", 3, 8, 'lower center')
    right = _add_legend(right)

    plt.show()

def print_objects(df_objects, base_point = None, print_id = True, title = None):
    """
    Визуализирует объекты в 3D-графике.

    Args:
        df_objects (pd.DataFrame): Данные об объектах.
        base_point (tuple, optional): Координаты базовой точки. По умолчанию None.
        print_id (bool, optional): Выводить идентификаторы объектов. По умолчанию True.
        title (str, optional): Заголовок подграфика. По умолчанию не задан.
    """
    ax = _prepare_ax()

    ax = _prepare_objects(ax, df_objects, print_id)

    if base_point is not None:
        ax = _prepare_base(ax, base_point)

    max_z = int(np.ceil(df_objects['z'].max() * 1.1 / 100) * 100)
    ax = _ax_end(ax, max_z, title)

    plt.show()


def print_clusters(df_objects, df_clusters, base_point = None):
    """
    Визуализирует кластеры объектов на 3D-графике.

    Args:
        df_objects (pd.DataFrame): Данные об объектах.
        df_clusters (pd.DataFrame): Данные о кластерах.
        base_point (tuple, optional): Координаты базовой станции. По умолчанию None.
    """
    ax = _prepare_ax()
    ax = _prepare_clusters(ax, df_objects, df_clusters)

    ax.set_title("Визуальное отображение всех связей главных объектов и базовой станции")

    if base_point is not None:
        ax = _prepare_base(ax, base_point)

    ax = _add_objects_ids(ax, df_objects, 'black', 'black')

    max_z = int(np.ceil(df_objects['z'].max() * 1.1 / 100) * 100)
    ax = _ax_end(ax, max_z, "Визуализация кластера и объектов на 3D-графике.")
    ax = _add_legend(ax)
    plt.show()


def print_leaders(df_objects, df_clusters, base_point = None, colors_line = None, show_objects = False):
    """
    Визуализирует главные объекты и маршруты связи на 3D-графике.

    Args:
        df_objects (pd.DataFrame): Данные об объектах.
        df_clusters (pd.DataFrame): Данные о кластерах.
        base_point (tuple, optional): Координаты базовой станции. По умолчанию None.
    """
    ax = _prepare_ax()

    ax = _prepare_clusters(ax, df_objects, df_clusters, 0.2)

    if base_point is not None:
        ax = _prepare_base(ax, base_point)

    ax = _add_objects_ids(ax, df_objects, 'black', 'black', show_objects)

    df_leaders = df_objects[df_objects['leader']]

    for i, leader_point1 in enumerate(df_leaders.itertuples()):
        for j, leader_point2 in enumerate(df_leaders.itertuples()):
            if i != j:
                line_width = 0.6
                line_style = "--"
                line_color = colors_line
                ax.plot([leader_point1.x, leader_point2.x],
                        [leader_point1.y, leader_point2.y],
                        [leader_point1.z, leader_point2.z],
                        color=line_color, linestyle=line_style, linewidth=line_width)

    for leader_point in df_leaders.itertuples():
        line_width = 0.6
        line_style = "--"
        line_color = colors_line
        legend_leader_connection, = ax.plot([base_point[0], leader_point.x],
                                            [base_point[1], leader_point.y],
                                            [base_point[2], leader_point.z],
                                            color=line_color, linestyle=line_style,
                                            linewidth=line_width)

    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('Высота')

    max_z = int(np.ceil(df_objects['z'].max() * 1.1 / 100) * 100)
    ax.set_zlim(0, max_z)

    handles, labels = ax.get_legend_handles_labels()
    handles.append(legend_leader_connection)
    labels.append("Маршруты связи")
    ax.set_title("Визуализация маршрутов связи кластеров")
    ax.legend(handles,labels, ncol=3, fontsize=8, loc="lower center")

    plt.show()

def print_graph_connections(Graph, title = None, show_objects = True, line_size=0.6):
    """
    Визуализирует связи графа, включая базовую станцию и объекты.

    Args:
        G (nx.Graph): Граф.
        title (str, optional): Заголовок графика. По умолчанию None.
        show_objects (bool, optional): Отображать объекты. По умолчанию True.
        line_size (float, optional): Толщина линий связей. По умолчанию 0.6.
    """
    ax = _prepare_ax()

    ax = _add_base_by_nodes(ax, Graph, 'base')

    if show_objects:
        ax = _add_objects_by_nodes(ax, Graph, ['object'], "Объект", None, True, FONT_SIZE, 'gray')

    ax = _add_objects_by_nodes(ax, Graph, ['leader'], "Главный кластера", 'black', True, FONT_SIZE_MEDIUM, 'black')

    if show_objects:
        ax = _add_edges_to_graph(ax, Graph, ['object'], False,'--', line_size, "Cвязь в кластере")

    ax = _add_edges_to_graph(ax, Graph, ['base_connector', 'base_connector_filtered'], 'g','--', line_size, "Cвязь c базовой станцией")
    ax = _add_edges_to_graph(ax, Graph, ['leaders','leaders_filtered'], False,'-', line_size, "Cвязь между кластерами")

    max_z = int((_get_max_z(Graph) * 1.1 / 100) * 100)
    _ax_end(ax, max_z, title)

    plt.show()

def print_short_paths(Graph, df_objects, df_clusters, paths, show_clusters = True):
    """
    Визуализирует кластеры объектов, маршруты связи и базовую станцию на 3D-графике.

    Args:
        Graph (nx.Graph): Граф.
        df_objects (pd.DataFrame): Данные об объектах.
        df_clusters (pd.DataFrame): Данные о кластерах.
        paths (list): Список начальных идентификаторов объектов для построения маршрутов.
    """
    ax = _prepare_ax()

    ax = _add_base_by_nodes(ax, Graph, 'base')

    if show_clusters:
        ax = _prepare_clusters(ax, df_objects, df_clusters, 0.2)
    else:
        ax = _prepare_objects(ax, df_objects, False)
        
    ax = _add_objects_ids(ax, df_objects, 'gray')

    ax = _add_edges_to_graph(ax, Graph, ['object'], False,'--', 0.6, "Cвязь в кластере")
    ax = _add_edges_to_graph(ax, Graph, ['base_connector'], 'g','--', 1, "Cвязь c базовой станцией")
    ax = _add_edges_to_graph(ax, Graph, ['leaders'], 'blue','--', 0.8, "Cвязь между кластерами")

    ax = _add_paths_to_graph(ax, Graph, df_objects, paths, 1.3)

    max_z = int((_get_max_z(Graph) * 1.1 / 100) * 100)
    _ax_end(ax, max_z, 'Визуализация кластеризации и маршрутов связи', 3, 8)

    plt.show()

def _prepare_objects(ax, df_objects, print_id = True, color = 'black', show_leaders = True, show_objects = True):
    """
    Подготавливает объекты для визуализации в 3D-графике.

    Args:
        ax (Axes3D): 3D-подграфик.
        df_objects (pd.DataFrame): Данные об объектах.
        print_id (bool, optional): Выводить идентификаторы объектов. По умолчанию True.
        color (str, optional): Цвет текста и точек объектов. По умолчанию 'black'.
        show_leaders (bool, optional): Выводить идентификаторы лидеров. По умолчанию True.
        show_objects (bool, optional): Выводить идентификаторы всех объектов. По умолчанию True.

    Returns:
        ax (Axes3D): Настроенный 3D-подграфик.
    """
    step_x = _get_step('x', df_objects['x'].max())
    step_y = _get_step('y', df_objects['y'].max())
    step_z = _get_step('z', df_objects['z'].max())

    if print_id is True:
        for _, obj in df_objects.iterrows():
            node = obj['object_id']

            if show_leaders and obj.get('leader', False):
                ax.text(obj['x'] + step_x,obj['y'] + step_y,obj['z'] + step_z, node , color=color, fontsize=FONT_SIZE_MEDIUM)
            elif show_objects:
                ax.text(obj['x'] + step_x,obj['y'] + step_y,obj['z'] + step_z, node , color=color, fontsize=FONT_SIZE)

        label = "Объект с идентификатором"
    else:
        label = "Объект"

    ax.scatter(df_objects['x'], df_objects['y'], df_objects['z'], label=label)

    return ax

def _prepare_clusters(ax, df_objects, df_clusters, alpha=0.4):
    """
    Подготавливает кластеры объектов для визуализации на 3D-графике.

    Args:
        ax (Axes3D): 3D-подграфик.
        df_objects (pd.DataFrame): Данные об объектах.
        df_clusters (pd.DataFrame): Данные о кластерах.
        alpha (float, optional): Прозрачность поверхности кластеров. По умолчанию 0.4.

    Returns:
        ax (Axes3D): Настроенный 3D-подграфик.
    """
    unique_clusters = df_objects['cluster_id'].unique()
    unique_clusters = np.sort(unique_clusters)

    colors = get_colors_palette(len(unique_clusters))
    dark_colors = colors * 0.8

    last_point = []
    for i, cluster_id in enumerate(unique_clusters):
        points = df_objects[df_objects['cluster_id'] == cluster_id]
        leader = points[points['leader']]
        leader_id = leader['object_id'].iloc[0]
        objects = points.loc[~points['leader']]

        last_point = {'color': dark_colors[i], 'x': leader['x'], 'y': leader['y'], 'z': leader['z']}
        label_text = f'Кластер ID: {cluster_id}, объектов: {len(points)}, главный ID: {leader_id}'
        ax.plot_trisurf(points['x'], points['y'], points['z'], color=colors[i], alpha=alpha, label=label_text)

        ax.scatter(objects['x'], objects['y'], objects['z'], c=[dark_colors[i]] * len(objects))

    ax.scatter(df_clusters['center_x'], df_clusters['center_y'], df_clusters['center_z'],  c="grey", s=40, label="Центроиды")

    cluster_leaders = df_objects[df_objects['leader']]
    ax.scatter(cluster_leaders['x'], cluster_leaders['y'], cluster_leaders['z'],  c="black", s=20, label="Главные объекты")

    # Добавим лейбл
    ax.scatter(last_point['x'], last_point['y'], last_point['z'], color=[last_point['color']], label="Объект")

    return ax

def _add_objects_ids(ax, df_objects, objects_color='gray', leaders_color='gray', show_objects=True, show_leaders=True):
    """
    Добавляет идентификаторы объектов на 3D-график.

    Args:
        ax (Axes3D): 3D-подграфик.
        df_objects (pd.DataFrame): Данные об объектах.
        objects_color (str, optional): Цвет текста для объектов. По умолчанию 'gray'.
        leaders_color (str, optional): Цвет текста для лидеров. По умолчанию 'gray'.
        show_objects (bool, optional): Выводить идентификаторы всех объектов. По умолчанию True.
        show_leaders (bool, optional): Выводить идентификаторы лидеров. По умолчанию True.

    Returns:
        ax (Axes3D): Настроенный 3D-подграфик.
    """
    step_x = _get_step('x', df_objects['x'].max())
    step_y = _get_step('y', df_objects['y'].max())
    step_z = _get_step('z', df_objects['z'].max())

    for _, obj in df_objects.iterrows():
        node = obj['object_id']

        if show_leaders and obj.get('leader', False):
            ax.text(obj['x'] + step_x,obj['y'] + step_y,obj['z'] + step_z, node , color=leaders_color, fontsize=FONT_SIZE_MEDIUM)
        elif show_objects:
            ax.text(obj['x'] +step_x,obj['y'] + step_y,obj['z'] + step_z, node , color=objects_color, fontsize=FONT_SIZE)

    return ax


def _add_paths_to_graph(ax, G, df_objects, paths, line_width = 1, line_color='black', line_style = "-"):
    """
    Добавляет маршруты связи на 3D-график.

    Args:
        ax (Axes3D): 3D-подграфик.
        G (nx.Graph): Граф.
        df_objects (pd.DataFrame): Данные об объектах.
        paths (list): Список начальных идентификаторов объектов для построения маршрутов.
        line_width (float, optional): Толщина линий маршрутов. По умолчанию 1.
        line_color (str, optional): Цвет линий маршрутов. По умолчанию 'black'.
        line_style (str, optional): Стиль линий маршрутов. По умолчанию '-'.

    Returns:
        ax (Axes3D): Настроенный 3D-подграфик.
    """
    pos = nx.get_node_attributes(G, 'pos')

    step_x = _get_step('x', df_objects['x'].max())
    step_y = _get_step('y', df_objects['y'].max())
    step_z = _get_step('z', df_objects['z'].max())

    step_text_xy = (df_objects['x'].max()) / 6
    step_text_z = (df_objects['z'].max()) / 20

    paths = _check_paths(df_objects, paths)

    for point, start_id in enumerate(paths):

        obj = df_objects[df_objects['object_id'] == start_id]

        shortest_path_value = obj['shortest_path'].iloc[0]

        for i in range(len(shortest_path_value) - 1):
            x1, y1, z1 = pos[shortest_path_value[i]]
            x2, y2, z2 = pos[shortest_path_value[i + 1]]
            ax.plot([x1, x2], [y1, y2], [z1, z2],
                    c=line_color,
                    linestyle=line_style,
                    linewidth=line_width)

            ax.text(x1 + step_x, y1 + step_y, z1 + step_z,
                    shortest_path_value[i],
                    color='black',
                    fontsize=FONT_SIZE_MEDIUM,
                    bbox=dict(facecolor='y', alpha=0.4))

        shortest_path_length = obj['shortest_path_length'].values[0]
        steps = obj['shortest_path_steps'].iloc[0]
        text = f"Маршрут от {start_id}, шагов: {steps}, длина пути: {shortest_path_length:.2f}"
        ax.text(step_text_xy, step_text_xy, 0 + point * step_text_z, text, color='black', fontsize=FONT_SIZE_MEDIUM)

    # Добавим легенду
    ax.plot([0,0], [0,0], [0,0], c=line_color, linestyle=line_style, linewidth=line_width, label="Маршрут связи")

    return ax

def _check_paths(df_objects, paths = False):
    """
    Если paths - это массив, возвращает paths.
    Если paths - это число, возвращает произвольные object_id из датасета not_leaders в количестве, равном paths.
    В противном случае возвращает пустой массив.

    Args:
        df_objects (pandas.DataFrame): Датасет с информацией об объектах.
        paths (bool, int, optional): Флаг или число. По умолчанию False.

    Returns:
        list: Массив object_id или пустой массив.
    """
    not_leaders = df_objects.loc[~df_objects['leader']]

    if paths == False:
        return []  # Возвращаем пустой массив
    elif isinstance(paths, list):
        return paths  # Возвращаем массив paths
    elif isinstance(paths, int):
        # Возвращаем произвольные object_id из датасета not_leaders в количестве, равном paths
        random_object_ids = not_leaders.sample(n=paths)['object_id'].tolist()
        return random_object_ids
    else:
        return []  # Возвращаем пустой массив

def _get_max_z(G):
    """
    Находит максимальное значение координаты Z в графе.

    Args:
        G (nx.Graph): Граф.

    Returns:
        max_z (float): Максимальное значение координаты Z.
    """
    pos = nx.get_node_attributes(G, 'pos')

    # Извлекаем значения координат Z
    z_values = [data[2] for _, data in pos.items()]

    # Находим максимальное значение
    max_z = max(z_values)

    return max_z

def _add_base_by_nodes(ax, G, group = 'base'):
    """
    Добавляет базовую станцию на 3D-график в соответствии с группой узлов.

    Args:
        ax (Axes3D): 3D-подграфик.
        G (nx.Graph): Граф.
        group (str, optional): Группа узлов для отображения базовой станции. По умолчанию 'base'.

    Returns:
        ax (Axes3D): Настроенный 3D-подграфик.
    """
    pos = nx.get_node_attributes(G, 'pos')

    for u, g in G.nodes(data=True):
        x, y, z = pos[u]
        if g.get('group') == group:  # Проверяем значение group
            ax = _prepare_base(ax, [x, y, z])

    return ax

def _add_objects_by_nodes(ax, G, group = 'all', obj_name = "Объект", object_color = None, print_id = True, obj_font_size = FONT_SIZE, id_color = None):
    """
    Добавляет объекты на 3D-график в соответствии с группой узлов.

    Args:
        ax (Axes3D): 3D-подграфик.
        G (nx.Graph): Граф.
        group (str, optional): Группа узлов для отображения объектов. По умолчанию 'all'.
        obj_name (str, optional): Название объектов. По умолчанию "Объект".
        object_color (str, optional): Цвет объектов. По умолчанию None.
        print_id (bool, optional): Выводить идентификаторы объектов. По умолчанию True.
        obj_font_size (int, optional): Размер шрифта для идентификаторов объектов. По умолчанию FONT_SIZE.
        id_color (str, optional): Цвет идентификаторов объектов. По умолчанию None.

    Returns:
        ax (Axes3D): Настроенный 3D-подграфик.
    """
    pos = nx.get_node_attributes(G, 'pos')

    obj = []
    for u, g in G.nodes(data=True):
        x, y, z = pos[u]
        if group == 'all' or g.get('group') in group: # Проверяем значение group
            obj.append({'node': u, 'x': x, 'y': y, 'z': z})

    obj_df = pd.DataFrame(obj)  # Преобразуем в DataFrame

    step_x = _get_step('x', obj_df['x'].max())
    step_y = _get_step('y', obj_df['y'].max())
    step_z = _get_step('z', obj_df['z'].max())

    if print_id is True:
        for _, obj in obj_df.iterrows():
            node = int(obj['node'])
            ax.text(obj['x'] + step_x,obj['y'] + step_y,obj['z'] + step_z, node , color=id_color, fontsize=obj_font_size)
        label = f"{obj_name} с идентификатором"
    else:
        label = f"{obj_name}"

    ax.scatter(obj_df['x'], obj_df['y'], obj_df['z'], color=object_color, label=label)

    return ax

def _get_step(coord = 'x', value = 10000):

    if value < 5000:
        step = 20
    else:
        if coord == 'x':
            step = STEP_X
        elif type == 'y':
            step = STEP_Y
        else:
            step = STEP_Z

    return step
    
def _add_edges_to_graph(ax, G, group = 'all', edge_color = None, line_style = '--', line_width = 1, label_text = False):
    """
    Добавляет ребра графа на 3D-график.

    Args:
        ax (Axes3D): 3D-подграфик.
        G (nx.Graph): Граф.
        group (str, optional): Группа ребер для отображения. По умолчанию 'all'.
        edge_color (str, optional): Цвет ребер. По умолчанию None.
        line_style (str, optional): Стиль линии ребер. По умолчанию '--'.
        line_width (float, optional): Толщина линии ребер. По умолчанию 1.
        label_text (bool, optional): Выводить лейбл ребер. По умолчанию False.

    Returns:
        ax (Axes3D): Настроенный 3D-подграфик.
    """
    pos = nx.get_node_attributes(G, 'pos')

    for u, v,  g in G.edges(data=True):
        x1, y1, z1 = pos[u]
        x2, y2, z2 = pos[v]

        if group == 'all' or g.get('group') in group:  # Проверяем значение group

            if edge_color:
                ax.plot([x1, x2], [y1, y2], [z1, z2], color=edge_color, linestyle=line_style, linewidth=line_width)
            else:
                ax.plot([x1, x2], [y1, y2], [z1, z2], linestyle=line_style, linewidth=line_width)

    # Добавим лейбл
    if label_text:
        if edge_color:
            ax.plot([0,0], [0,0], [0,0], color=edge_color, linestyle=line_style, linewidth=line_width, label=label_text)
        else:
            ax.plot([0,0], [0,0], [0,0], linestyle=line_style, linewidth=line_width, label=label_text)

    return ax

def _prepare_base(ax, base_point):
    """
    Подготавливает базовую станцию для визуализации на 3D-графике.

    Args:
        ax (Axes3D): 3D-подграфик.
        base_point (tuple): Координаты базовой станции.

    Returns:
        ax (Axes3D): Настроенный 3D-подграфик.
    """

    step_z = 20
    step_x = 20
    step_y = 20

    ax.scatter(base_point[0], base_point[1], base_point[2], color='red', label="Базовая станция")
    ax.text(base_point[0] + step_x, base_point[1] + step_y, base_point[2] - step_z, 'Базовая станция', color='red', fontsize=FONT_SIZE_MEDIUM)

    return ax

def _prepare_axs():
    """
    Создает два 3D-подграфика (один слева, другой справа) в одном ряду.

    Returns:
        left (Axes3D): Подграфик слева.
        right (Axes3D): Подграфик справа.
    """
    _, (left, right) = plt.subplots(1, 2, figsize=FIG_SIZE, subplot_kw={'projection': '3d'})

    return left, right

def _prepare_ax(title = False):
    """
    Создает 3D-подграфик с возможностью добавления заголовка.

    Args:
        title (str, optional): Заголовок подграфика. По умолчанию не задан.

    Returns:
        ax (Axes3D): 3D-подграфик.
    """
    fig = plt.figure(figsize=FIG_SIZE)
    ax = fig.add_subplot(111, projection='3d')

    if title:
        ax.set_title(title)

    return ax

def _ax_end(ax, max_z = 1200, title=None, legend_col = 1, legend_font=8, legend_loc='lower center'):
    """
    Настраивает параметры 3D-подграфика для завершения визуализации.

    Args:
        ax (Axes3D): 3D-подграфик.
        max_z (float, optional): Максимальное значение по оси Z. По умолчанию 1200.
        title (str, optional): Заголовок подграфика. По умолчанию не задан.
        legend_col (int, optional): Количество столбцов в легенде. По умолчанию 1.
        legend_font (int, optional): Размер шрифта в легенде. По умолчанию 8.
        legend_loc (str, optional): Расположение легенды. По умолчанию 'lower right'.

    Returns:
        ax (Axes3D): Настроенный 3D-подграфик.
    """
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Высота')

    if title:
        ax.set_title(title)

    ax.set_zlim(0, max_z)

    ax.legend(scatterpoints=1,
            loc=legend_loc,
            ncol=legend_col,
            fontsize=legend_font)

    return ax

def _add_legend(ax, legend_loc = 'lower center' , legend_col = 3, legend_font = 8):
    """
    Добавляет легенду на 3D-график.

    Args:
        ax (Axes3D): 3D-подграфик.
        legend_loc (str, optional): Расположение легенды. По умолчанию 'lower right'.
        legend_col (int, optional): Количество столбцов в легенде. По умолчанию 3.
        legend_font (int, optional): Размер шрифта в легенде. По умолчанию 8.

    Returns:
        ax (Axes3D): Настроенный 3D-подграфик.
    """
    ax.legend(scatterpoints=1,
            loc=legend_loc,
            ncol=legend_col,
            fontsize=legend_font)

    return ax

def get_colors_list(count):
    """
    Возвращает список случайных цветов из заданного набора.

    Args:
        count (int): Количество цветов для возврата.

    Returns:
        list: Список случайных цветов.
    """
    color_names = ['red', 'green', 'blue', 'orange', 'purple', 'gray']
    extended_color_names = np.random.choice(color_names, size=count, replace=True)

    return extended_color_names

def get_colors_palette(count):
    """
    Возвращает палитру цветов из заданного набора.

    Args:
        count (int): Количество цветов в палитре.

    Returns:
        np.ndarray: Палитра цветов.
    """
    viridis = plt.colormaps['tab20'].resampled(count)
    colors = viridis(np.linspace(0, 1, count))

    return colors
